import os

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

os.environ['LANGCHAIN_TRACING_V2'] = "true"
os.environ['LANGCHAIN_API_KEY'] = '1123'

# 调用大模型

model = ChatOpenAI(model='gpt-4-turbo', temperature=0)


# class Classification(BaseModel):
#     """
# 定义一个pydantic模型，未来需要根据该类型，完成文本的分类
#     """
#     sentiment: str = Field(description="文本的情感")
#     aggressiveness: int = Field(description="描述文本的攻击性，数字越大表示越攻击性")
#     language: str = Field(description="文本使用的语言")


class Classification(BaseModel):
    """
定义一个pydantic模型，未来需要根据该类型，完成文本的分类
    """
    sentiment: str = Field(..., enum=["happy", "neutral", "sad"], description="文本的情感")
    aggressiveness: int = Field(..., enum=[1, 2, 3, 4, 5], description="描述文本的攻击性，数字越大表示越攻击性")
    language: str = Field(..., enum=["english", "chinese"], description="文本使用的语言")


# 创建一个用于提取新的prompt
tagging_prompt = ChatPromptTemplate.from_template(
    """
    从以下段落中提取所需信息，只提取‘Classification’类中提到的属性
    段落：
    {input}
    """
)
chain = tagging_prompt | model.with_structured_output(Classification)
resp = chain.invoke({'input': '道德败坏，做出的事情令人生气'})
